package com.java.sample.encrypt.ecdh;

import com.java.sample.common.util.Base64Util;

import javax.crypto.Cipher;
import javax.crypto.KeyAgreement;
import javax.crypto.SecretKey;
import javax.crypto.spec.GCMParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.security.*;
import java.security.spec.ECGenParameterSpec;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;

public class EcdhUtils {

    /**
     * 生成密钥配置信息
     * @return
     */
    public static EcdhKeyConfig genKeyConfig(){
        KeyPair keyPair = generateKeyPair();
        EcdhKeyConfig keyConfig = new EcdhKeyConfig();
        keyConfig.setPrivateKey(Base64Util.encode(keyPair.getPrivate().getEncoded()));
        keyConfig.setPublicKey(Base64Util.encode(keyPair.getPublic().getEncoded()));
        keyConfig.setUncompressedPublicKey(WebEcdhUtils.getUncompressedPublicKey(keyPair.getPublic()));
        return keyConfig;
    }

    /**
     * 计算共享密钥
     * @param privateKey
     * @param publicKey
     * @return
     */
    public static byte[] calculateSharedSecret(PrivateKey privateKey, PublicKey publicKey) {
        try {
            KeyAgreement keyAgreement = KeyAgreement.getInstance("ECDH");
            keyAgreement.init(privateKey);
            keyAgreement.doPhase(publicKey, true);
            return keyAgreement.generateSecret();
        } catch (Exception e){
            e.printStackTrace();
            throw new RuntimeException("calculateSharedSecret error");
        }
    }

    /**
     * 派生AES密钥（使用SHA-256处理共享密钥）
     * @param sharedSecret
     * @return
     */
    public static SecretKey deriveAesKey(byte[] sharedSecret) {
        try {
            MessageDigest digest = MessageDigest.getInstance("SHA-256");
            byte[] aesKeyBytes = digest.digest(sharedSecret);
            return new SecretKeySpec(aesKeyBytes, "AES");
        } catch (Exception e){
            e.printStackTrace();
            throw new RuntimeException("deriveAesKey error");
        }
    }

    /**
     * 加密消息
     * @param message
     * @param aesKey
     * @return
     */
    public static String encrypt(String message, SecretKey aesKey) {
        try{
            Cipher cipher = Cipher.getInstance("AES/GCM/NoPadding");
            byte[] iv = new byte[12]; // GCM推荐12字节IV
            SecureRandom random = new SecureRandom();
            random.nextBytes(iv);
            GCMParameterSpec parameterSpec = new GCMParameterSpec(128, iv);
            cipher.init(Cipher.ENCRYPT_MODE, aesKey, parameterSpec);
            byte[] encryptedData = cipher.doFinal(message.getBytes("UTF-8"));

            // 组合IV和加密数据
            byte[] combined = new byte[iv.length + encryptedData.length];
            System.arraycopy(iv, 0, combined, 0, iv.length);
            System.arraycopy(encryptedData, 0, combined, iv.length, encryptedData.length);

            return Base64.getEncoder().encodeToString(combined);
        } catch (Exception e){
            e.printStackTrace();
            throw new RuntimeException("encrypt error");
        }
    }

    /**
     * 解密消息
     * @param encryptedMessage
     * @param aesKey
     * @return
     */
    public static String decrypt(String encryptedMessage, SecretKey aesKey){
        try {
            byte[] combined = Base64.getDecoder().decode(encryptedMessage);
            byte[] iv = new byte[12];
            byte[] encryptedData = new byte[combined.length - 12];

            System.arraycopy(combined, 0, iv, 0, 12);
            System.arraycopy(combined, 12, encryptedData, 0, encryptedData.length);

            Cipher cipher = Cipher.getInstance("AES/GCM/NoPadding");
            GCMParameterSpec parameterSpec = new GCMParameterSpec(128, iv);
            cipher.init(Cipher.DECRYPT_MODE, aesKey, parameterSpec);
            byte[] decryptedData = cipher.doFinal(encryptedData);

            return new String(decryptedData, "UTF-8");
        }  catch (Exception e){
            e.printStackTrace();
            throw new RuntimeException("decrypt error");
        }
    }

    /**
     * 从Base64字符串加载私钥
     * @param privateKeyBase64
     * @return
     */
    public static PrivateKey loadPrivateKeyFromBase64(String privateKeyBase64) {
        try{
            byte[] decodedKey = Base64.getDecoder().decode(privateKeyBase64);
            PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(decodedKey);
            KeyFactory keyFactory = KeyFactory.getInstance("EC");
            return keyFactory.generatePrivate(keySpec);
        } catch (Exception e){
            e.printStackTrace();
            throw new RuntimeException("loadPrivateKeyFromBase64 error");
        }
    }

    /**
     * 从Base64字符串加载公钥
     * @param publicKeyBase64
     * @return
     */
    public static PublicKey loadPublicKeyFromBase64(String publicKeyBase64) {
        try{
            byte[] decodedKey = Base64.getDecoder().decode(publicKeyBase64);
            X509EncodedKeySpec keySpec = new X509EncodedKeySpec(decodedKey);
            KeyFactory keyFactory = KeyFactory.getInstance("EC");
            return keyFactory.generatePublic(keySpec);
        } catch (Exception e){
            e.printStackTrace();
            throw new RuntimeException("loadPublicKeyFromBase64 error");
        }
    }
    /**
     * 生成ECDH密钥对
     * @return
     */
    public static KeyPair generateKeyPair() {
        try{
            KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("EC");
            ECGenParameterSpec ecSpec = new ECGenParameterSpec("secp256r1");
            keyPairGenerator.initialize(ecSpec);
            return keyPairGenerator.generateKeyPair();
        } catch (Exception e) {
            e.printStackTrace();
            throw new RuntimeException("generateKeyPair error");
        }
    }

}
